import os
import sys
import json
sys.path.append('./')
import argparse

import config as _C
from utils.subprocess_utils import run_command_with_conda

conda_env = "dgl" if os.getenv("USER") == "soyongs" else "fbcontact"

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--sequence', default='')
    args = parser.parse_args()

    if args.sequence != '':
        _C.SEQUENCE_NAME = args.sequence
    
    metadata = os.path.join(_C.RAW_DATA_DIR, _C.SEQUENCE_NAME, "proc.json")
    with open(metadata, "rb") as f:
        proc_info = json.load(f)

    gender = proc_info["gender"]
    bbox_dir = os.path.join(_C.BBOX_DIR, _C.SEQUENCE_NAME)
    calib_pth = os.path.join(_C.COLMAP_DIR, _C.SEQUENCE_NAME, "calib.npz")
    vitpose_results_dir = os.path.join(_C.VITPOSE_RESULTS_DIR, _C.SEQUENCE_NAME)
    dense_vitpose_results_dir = os.path.join(_C.DENSE_VITPOSE_RESULTS_DIR, _C.SEQUENCE_NAME)
    smplifyx_results_pth = os.path.join(_C.SMPLIFYX_DIR, _C.SEQUENCE_NAME, "results.pkl")
    bbox_removal_pth = os.path.join(bbox_dir, "bbox_removal.json")
    large_regularization_json_pth = os.path.join(os.path.dirname(smplifyx_results_pth), "high_regularization_frames.json")
    os.makedirs(os.path.dirname(smplifyx_results_pth), exist_ok=True)
    
    working_dir = _C.SMPLIFYX_DIR
    run_cmd = [
        "python", "-u", 
        "-m", "run_fitting", 
        "--vitpose_results_dir", vitpose_results_dir, 
        "--dense_vitpose_results_dir", dense_vitpose_results_dir, 
        "--calib_path", calib_pth,
        "--bbox_dir", bbox_dir,
        "--output_results_pth", smplifyx_results_pth,
        "--gender", gender,
        "--vitpose_ignore_json_pth", bbox_removal_pth,
        "--large_regularization_json_pth", large_regularization_json_pth,
    ]

    run_command_with_conda(working_dir, conda_env, run_cmd)